from abc import ABC, abstractmethod
import os
import math
from datetime import datetime
import traceback

from loguru import logger
from astropy.io import fits
import numpy as np

from tools.base_tool import psf_comd, sex_comd
from tools.base_tool import get_filename_without_extension

class PSFPhotometry:
    def __init__(self, image_path, rootpath=None,
                 detect_minarea=5, detect_thresh=1.2, analysis_thresh=1.2):
        """
        初始化 PSFPhotometry 类。

        参数:
            image_path (str): 输入的 FITS 图像文件路径。
            rootpath (str, 可选): 根目录路径；默认为当前工作目录。
            detect_minarea (int, 可选): SExtractor 的 DETECT_MINAREA 参数。
            detect_thresh (float, 可选): SExtractor 的 DETECT_THRESH 参数。
            analysis_thresh (float, 可选): SExtractor 的 ANALYSIS_THRESH 参数。
        """
        self.image_path = image_path
        self.rootpath = rootpath or os.getcwd()

        # 目录和配置文件路径
        self.configdir = os.path.join(self.rootpath, "config")
        self.sexConf = os.path.join(self.configdir, "default.sex")
        self.sexParam = os.path.join(self.configdir, "default.paramNew")
        self.sexConfWithoutPsf = os.path.join(self.configdir, "default_ori.sex")
        self.sexParamWithoutPsf = os.path.join(self.configdir, "default_psfex.param")
        self.psfexConf = os.path.join(self.configdir, "default.psfex")
        self.figdir = os.path.join(self.rootpath, "figdir")

        
        # 生成产品文件的路径
        self.ldac_file = self._replace_extension('.ldac',type_="_psf")
        self.psf_file = self._replace_extension('.psf',type_="_psf")
        
        self.filename = get_filename_without_extension(self.image_path)

        # 初始化头信息
        self.init_header()

        # 初始化 SExtractor 命令模板，传入可配置的探测阈值
        self.init_sex_command(detect_minarea, detect_thresh, analysis_thresh)

        self.init_sex_command_wothout(detect_minarea, detect_thresh, analysis_thresh)
    def _replace_extension(self, new_ext,type_="_psf"):
        """替换文件扩展名"""
        return os.path.splitext(self.image_path)[0]+type_ + new_ext

    def init_header(self):
        """初始化 FITS 头信息"""
        hdr = fits.getheader(self.image_path)
        self.filterid = hdr.get("FILTER", "UNKNOWN")
        self.exptime = hdr.get("EXPOSURE", 1)
        self.objID = hdr.get("OBJECT", "UNKNOWN")
        self.date_ = hdr.get("DATE", "1970-01-01T00:00:00")
        self.saturate_value = hdr.get("SATURATE", 48000)

        # 解析并格式化日期
        try:
            dt = datetime.strptime(self.date_, '%Y-%m-%dT%H:%M:%S')
            self.formatted_date = dt.strftime('%Y%m%d')
        except ValueError:
            logger.warning(f"头信息中的日期格式不正确: {self.date_}")
            self.formatted_date = "19700101"

    def init_sex_command(self, detect_minarea, detect_thresh, analysis_thresh):
        """初始化 SExtractor 命令模板，将探测阈值作为参数传入"""
        self.sex_command_template = (
            "sex {image} -c {config} "
            "-PARAMETERS_NAME {param} "
            "-CATALOG_NAME {ldac} "
            "-PSF_NAME {psf} "
            f"-DETECT_MINAREA {detect_minarea} "
            f"-DETECT_THRESH {detect_thresh} "
            f"-ANALYSIS_THRESH {analysis_thresh} "
            f"-SATUR_LEVEL {self.saturate_value} "
            "-CHECKIMAGE_TYPE NONE -VERBOSE_TYPE NORMAL"
        )
        
    def init_sex_command_wothout(self, detect_minarea, detect_thresh, analysis_thresh):
        """初始化 SExtractor 命令模板，将探测阈值作为参数传入"""
        self.sex_command_template_without_psf = (
            "sex {image} -c {config} "
            "-PARAMETERS_NAME {param} "
            "-CATALOG_NAME {ldac} "
            f"-DETECT_MINAREA {detect_minarea} "
            f"-DETECT_THRESH {detect_thresh} "
            f"-ANALYSIS_THRESH {analysis_thresh} "
            f"-SATUR_LEVEL {self.saturate_value} "
            "-CHECKIMAGE_TYPE NONE -VERBOSE_TYPE NORMAL"
        )
        

    def run_sextractor_withoutpsf(self):
      
        # 格式化 SExtractor 命令
        sex_command = self.sex_command_template_without_psf.format(
            image=self.image_path,
            config=self.sexConfWithoutPsf,
            param=self.sexParamWithoutPsf,
            ldac=self.ldac_file
        )

        logger.debug(f"SExtractor PSF 命令: {sex_command}")
        sex_comd(sex_command)

        # 检查 LDAC 文件是否生成
        if not os.path.exists(self.ldac_file):
            logger.error(f"SExtractor 未能创建 LDAC 文件: {self.ldac_file}")
            raise FileNotFoundError(f"未找到 LDAC 文件: {self.ldac_file}")
        logger.info(f"SExtractor 测光完成: {self.ldac_file}")
    
    def run_sextractor_psf(self):
        """运行 SExtractor 进行 PSF 测光"""
        logger.info("运行 SExtractor 进行 PSF 测光。")

        # 格式化 SExtractor 命令
        sex_command = self.sex_command_template.format(
            image=self.image_path,
            config=self.sexConf,
            param=self.sexParam,
            ldac=self.ldac_file,
            psf=self.psf_file
        )

        logger.debug(f"SExtractor PSF 命令: {sex_command}")
        sex_comd(sex_command)

        # 检查 LDAC 文件是否生成
        if not os.path.exists(self.ldac_file):
            logger.error(f"SExtractor 未能创建 LDAC 文件: {self.ldac_file}")
            raise FileNotFoundError(f"未找到 LDAC 文件: {self.ldac_file}")
        logger.info(f"SExtractor PSF 测光完成: {self.ldac_file}")

    def run_psfex(self):
        """使用 PSFEx 生成 PSF 模型"""
        logger.info("使用 PSFEx 生成 PSF 模型。")

        # 打开 LDAC 文件并筛选源
        with fits.open(self.ldac_file) as hdul:
            sources = hdul[2].data
        clean_sources = sources[sources["FLAGS"] == 0]

        # 估算 FWHM 和内核大小
        fwhm_median = np.nanmedian(clean_sources["FWHM_IMAGE"])
        kernel_size = min(int(fwhm_median) * 7, 42)
        fwhm_range = '2.0,10.0' if fwhm_median <= 9 else '5.0,18.0'

        logger.info(f"估算的 FWHM: {fwhm_median}, 内核大小: {kernel_size}")

        # 定义 PSFEx 命令
        checkimage_dir = os.path.join(os.path.dirname(self.ldac_file), 'psf_checkimages')
        os.makedirs(checkimage_dir, exist_ok=True)

        psfex_command = (
            f"psfex {self.ldac_file} -c {self.psfexConf} "
            f"-PSF_SIZE {kernel_size},{kernel_size} "
            f"-PSF_DIR {os.path.dirname(self.ldac_file)} "
            f"-SAMPLE_FWHMRANGE {fwhm_range} "
            f"-CHECKPLOT_TYPE FWHM,ELLIPTICITY,COUNTS,CHI2,RESIDUALS "
            f"-CHECKPLOT_NAME {checkimage_dir}/fwhm,{checkimage_dir}/ellipticity,{checkimage_dir}/counts,{checkimage_dir}/chi2,{checkimage_dir}/resi "
            
            f"-CHECKIMAGE_TYPE CHI,PROTOTYPES,SAMPLES,RESIDUALS,SNAPSHOTS "
            f"-CHECKIMAGE_NAME {','.join([os.path.join(checkimage_dir, name) for name in ['chi.fits', 'proto.fits', 'samp.fits', 'resi.fits', 'snap.fits']])} "
            f"-XML_NAME {os.path.join(checkimage_dir, 'psfex.xml')}"
        )
        
        
        logger.debug(f"PSFEx 命令: {psfex_command}")
        psf_comd(psfex_command)

        # 检查 PSF 文件是否生成
        if not os.path.exists(self.psf_file):
            logger.error(f"PSFEx 未能创建 PSF 文件: {self.psf_file}")
            raise FileNotFoundError(f"未找到 PSF 文件: {self.psf_file}")
        logger.info(f"PSF 模型生成完成: {self.psf_file}")

    
    def execute(self):
        """执行 PSF 测光的所有步骤"""
        try:
            
            self.run_sextractor_withoutpsf()
            
            self.run_psfex()
            
            self.run_sextractor_psf()
            
            
            logger.info("PSF 测光工作流程成功完成。")
        except Exception as e:
            logger.error(f"PSF 测光过程中发生错误: {e}")
            logger.debug(traceback.format_exc())
            raise
 